
from torch.utils.tensorboard import SummaryWriter
from utils import sys_logger 
from models.base_model import BaseModel
from config import config_singleton
from model_manage import ModelManager
import time
from tcp_handle import TcpClient
from utils import Constants
import struct
from utils import BaseClient

class Collector():
    """
    """
    def __init__(self, writetoboard:bool, base_path, manager:ModelManager, handler:BaseClient):

        self.is_writetoboard = writetoboard
        self.log_path = base_path
        self.models_list = manager.get_all_models()
        self.need_stop = False

        if self.is_writetoboard:
            self.summaryWriter = SummaryWriter(config_singleton.tensorboard_log_root)
        else:
            self.summaryWriter = None
        
        self.step_count = 0
        self.collect_cnt = 0
        self.client_handler = handler

        # [0xAFA0A0AF, frame_len, func_no, frame_num, cmd, 0x8100, timestamp==0, 0x0000, t]
        # 默认frmae_len= 0, func_no = 0x01, status = 0x8100, timestamp = 0,
        self.frame_head = [Constants.TCP_FRAME_HEAD, 0, 0x01, 0, 0x00, 0x8100, 0x00, 0x00, 0]
        self.frame_tail = [0x00, Constants.TCP_FRAME_TAIL]          # 默认将和校验填写为0x00
        self.tial_bytes = struct.pack("<HI", *self.frame_tail)


    def do_collect(self):
        """
        数据收集函数
        """
        while not self.need_stop:
            # 防止线程占用率较高，每50个循环切换一次
            self.step_count += 1
            if self.step_count % 50 == 0:
                time.sleep(0.1)

            for models in self.models_list:
                for model in models:
                    if isinstance(model, BaseModel):
                        vals = model.get_output_data()
                        # 判断数据部分是否都是None
                        if vals[0] is not None:     
                            self.collect_cnt += 1

                            # 如果需要写入tensorboard
                            # 
                            if self.is_writetoboard and self.summaryWriter != None:
                                varnames = model.get_output_varnames()
                                # model_name = model.get_model_name()
                                # all_val = {varname : val for varname, val in zip(varnames, vals[1])}
                                for varname, val in zip(varnames, vals[1]):
                                    # vals[0] timestamp 最小分辨率为0.001s
                                    if varname.startswith("@"):
                                        self.summaryWriter.add_scalar(varname, val, vals[0] * 1000)
                                # self.summaryWriter.add_scalars(model_name, all_val, vals[0] * 1000)
                            
                            # 将数据发送到TCP服务器
                            byte_array = self.msg_encoder(model.output_table_num, vals)

                            if byte_array is None:
                                continue
                            # 添加发送数据to the TCP client
                            self.client_handler.put_send_data(byte_array)
                            

        
        sys_logger.info(f"Collector stopped collecting data. total collected {self.collect_cnt} frame data.")

    
    def stop(self):
        """
        关闭线程，关闭写入
        """
        if self.summaryWriter is not None:
            self.summaryWriter.close()
            self.summaryWriter = None

        self.need_stop = True



    def msg_encoder(self, table_num, vals):
        """
        协议编码函数
        """
        if vals[0] is None or vals[1] is None:
            return None

        float_cnt = len(vals[1])        # 第一个为时间

        self.frame_head[1] = 33 + 4 * float_cnt                # frame len
        self.frame_head[3] = self.collect_cnt                  # frame cnt
        self.frame_head[4] = table_num                         # cmd word
        # self.frame_head[8] = int(vals[0] * 1000)               # t 转化为ms，并转化为无符号整形
        self.frame_head[8] = vals[0]

        # temp_str = f'<IIBIHHIHI{float_cnt}f'
        byte_fromate = f'<IIBIHHIHf{float_cnt}f'
        

        if type(vals[1]) == list:
            byte_array = struct.pack(byte_fromate, *(self.frame_head + vals[1]))
        else:
            byte_array = struct.pack(byte_fromate, *(self.frame_head + vals[1].tolist()))

        byte_array = byte_array + self.tial_bytes           # TODO：默认将和校验填写为0x00， 后续增加校验功能

        # print(data.hex())
        # sys_logger.info(byte_array.hex())
        # sys_logger.info(f"\t{table_num}\t{vals[0]}\t{byte_array.hex()}")

        return byte_array 








